#!/usr/bin/python3
# -*- coding: utf-8 -*-
import codecs
import copy
import getopt
import hashlib
import json
#import os
import re
#import sys
import time
import urllib
from threading import Thread

import requests
from six.moves import queue as Queue

import sys
import os
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import *

# Setting timeout
TIMEOUT = 10

# Retry times
RETRY = 5

# Numbers of downloading threads concurrently
THREADS = 2

HEADERS = {
    'accept-encoding': 'gzip, deflate, br',
    'accept-language': 'zh-CN,zh;q=0.9',
    'pragma': 'no-cache',
    'cache-control': 'no-cache',
    'upgrade-insecure-requests': '1',
    'user-agent': "Mozilla/5.0 (iPhone; CPU iPhone OS 11_0 like Mac OS X) AppleWebKit/604.1.38 (KHTML, like Gecko) Version/11.0 Mobile/15A372 Safari/604.1",
}

def download(medium_type, uri, medium_url, target_folder):

    headers = copy.copy(HEADERS)
    file_name = uri
    if medium_type == 'video':
        file_name += '.mp4'
        headers['user-agent'] = 'Aweme/27014 CFNetwork/974.2.1 Darwin/18.0.0'
    elif medium_type == 'image':
        file_name += '.jpg'
        file_name = file_name.replace("/", "-")
    else:
        return

    file_path = os.path.join(target_folder, file_name)
    if os.path.isfile(file_path):
        TextWrite(file_name + " 已经爬取过了，文件保存在 " + file_path + " 放弃爬取")
        return

    TextWrite("Downloading %s from %s.\n" % (file_name, medium_url))
    # VIDEOID_DICT[VIDEO_ID] = 1  # 记录已经下载的视频
    retry_times = 0
    while retry_times < RETRY:
        try:
            resp = requests.get(medium_url, headers=headers, stream=True, timeout=TIMEOUT)
            if resp.status_code == 403:
                retry_times = RETRY
                TextWrite("Access Denied when retrieve %s.\n" % medium_url)
                raise Exception("Access Denied")
            with open(file_path, 'wb') as fh:
                for chunk in resp.iter_content(chunk_size=1024):
                    fh.write(chunk)
            break
        except:
            pass
        
        retry_times += 1
    else:
        try:
            os.remove(file_path)
        except OSError:
            pass
        TextWrite("Failed to retrieve %s from %s.\n" % (uri, medium_url))
        time.sleep(1)

def get_real_address(url):
    if url.find('v.douyin.com') < 0:
        return url
    res = requests.get(url, headers=HEADERS, allow_redirects=False)
    return res.headers['Location'] if res.status_code == 302 else None

def get_dytk(url):
    res = requests.get(url, headers=HEADERS)
    if not res:
        return None
    dytk = re.findall("dytk: '(.*)'", res.content.decode('utf-8'))
    if len(dytk):
        return dytk[0]
    return None

class DownloadWorker(Thread):
    def __init__(self, queue):
        Thread.__init__(self)
        self.queue = queue

    def run(self):
        while True:
            medium_type, uri, download_url, target_folder = self.queue.get()
            download(medium_type, uri, download_url, target_folder)
            self.queue.task_done()

class CrawlerScheduler(object):

    def __init__(self, items):
        self.numbers = []
        self.challenges = []
        self.musics = []
        for i in range(len(items)):
            url = get_real_address(items[i])
            if not url:
                continue
            if re.search('share/user', url):
                self.numbers.append(url)
            if re.search('share/challenge', url):
                self.challenges.append(url)
            if re.search('share/music', url):
                self.musics.append(url)

        self.queue = Queue.Queue()
        self.scheduling()

    @staticmethod
    def generateSignature(value):
        p = os.popen('node acrawler.js %s' % value) #需要安装好nodejs
        return p.readlines()[0]

    @staticmethod
    def calculateFileMd5(filename):
        hmd5 = hashlib.md5()
        fp = open(filename, "rb")
        hmd5.update(fp.read())
        return hmd5.hexdigest()

    def scheduling(self):
        for x in range(THREADS):
            worker = DownloadWorker(self.queue)
            worker.daemon = True
            worker.start()

        for url in self.numbers:
            self.download_user_videos(url)
        for url in self.challenges:
            self.download_challenge_videos(url)
        for url in self.musics:
            self.download_music_videos(url)

    def download_user_videos(self, url):
        number = re.findall(r'share/user/(\d+)', url)
        if not len(number):
            return
        dytk = get_dytk(url)
        hostname = urllib.parse.urlparse(url).hostname
        if hostname != 't.tiktok.com' and not dytk:
            return
        user_id = number[0]
        video_count = self._download_user_media(user_id, dytk, url)
        self.queue.join()
        TextWrite("\nAweme number %s, video number %s\n\n" %
              (user_id, str(video_count)))
        TextWrite("\nFinish Downloading All the videos from %s\n\n" % user_id)

    def download_challenge_videos(self, url):
        challenge = re.findall('share/challenge/(\d+)', url)
        if not len(challenge):
            return
        challenges_id = challenge[0]
        video_count = self._download_challenge_media(challenges_id, url)
        self.queue.join()
        TextWrite("\nAweme challenge #%s, video number %d\n\n" %
              (challenges_id, video_count))
        TextWrite("\nFinish Downloading All the videos from #%s\n\n" % challenges_id)


    def _join_download_queue(self, aweme, target_folder):
        try:
            if aweme.get('video', None):
                uri = aweme['video']['play_addr']['uri']
                download_url = "https://aweme.snssdk.com/aweme/v1/play/?{0}" #无水印
                download_params = {
                    'video_id': uri,
                    'line': '0',
                    'ratio': '720p',
                    'media_type': '4',
                    'vr_type': '0',
                    'improve_bitrate': '0',
                    'is_play_url' : '1',
                    'is_support_h265' : '0',
                    'source' : 'PackSourceEnum_PUBLISH'
                }


                share_info = aweme.get('share_info', {})
                url = download_url.format(
                    '&'.join([key + '=' + download_params[key] for key in download_params]))
                self.queue.put(('video',
                                uri + "-" + share_info.get('share_desc', uri),
                                url, target_folder))
            else:
                if aweme.get('image_infos', None):
                    image = aweme['image_infos']['label_large']
                    self.queue.put(
                        ('image', image['uri'], image['url_list'][0], target_folder))

        except KeyError:
            return
        except UnicodeDecodeError:
            TextWrite("Cannot decode response data from DESC %s" % aweme['desc'])
            return

    def _download_user_media(self, user_id, dytk, url):
        current_folder = os.getcwd()
        target_folder = os.path.join(current_folder, 'download/%s' % user_id)
        if not os.path.isdir(target_folder):
            os.mkdir(target_folder)

        if not user_id:
            TextWrite("Number %s does not exist" % user_id)
            return
        hostname = urllib.parse.urlparse(url).hostname
        signature = self.generateSignature(str(user_id))
        user_video_url = "https://%s/aweme/v1/aweme/post/" % hostname
        user_video_params = {
            'user_id': str(user_id),
            'count': '21',
            'max_cursor': '0',
            'aid': '1128',
            '_signature': signature,
            'dytk': dytk
        }

        max_cursor, video_count = None, 0
        while True:
            if max_cursor:
                user_video_params['max_cursor'] = str(max_cursor)
            res = requests.get(user_video_url, headers=HEADERS,
                               params=user_video_params)
            contentJson = json.loads(res.content.decode('utf-8'))
            aweme_list = contentJson.get('aweme_list', [])
            for aweme in aweme_list:
                video_count += 1
                aweme['hostname'] = hostname
                self._join_download_queue(aweme, target_folder)
            if contentJson.get('has_more'):
                max_cursor = contentJson.get('max_cursor')
            else:
                break

        if video_count == 0:
            TextWrite("There's no video in number %s." % user_id)

        return video_count

    def _download_challenge_media(self, challenge_id, url):
        if not challenge_id:
            TextWrite("Challenge #%s does not exist" % challenge_id)
            return
        current_folder = os.getcwd()
        target_folder = os.path.join(
            current_folder, 'download/#%s' % challenge_id)
        if not os.path.isdir(target_folder):
            os.mkdir(target_folder)

        hostname = urllib.parse.urlparse(url).hostname
        signature = self.generateSignature(str(challenge_id) + '9' + '0')

        challenge_video_url = "https://%s/aweme/v1/challenge/aweme/" % hostname
        challenge_video_params = {
            'ch_id': str(challenge_id),
            'count': '9',
            'cursor': '0',
            'aid': '1128',
            'screen_limit': '3',
            'download_click_limit': '0',
            '_signature': signature
        }

        cursor, video_count = None, 0
        while True:
            if cursor:
                challenge_video_params['cursor'] = str(cursor)
                challenge_video_params['_signature'] = self.generateSignature(
                    str(challenge_id) + '9' + str(cursor))
            res = requests.get(challenge_video_url,
                               headers=HEADERS, params=challenge_video_params)
            try:
                contentJson = json.loads(res.content.decode('utf-8'))
            except:
                TextWrite(res.content)
            aweme_list = contentJson.get('aweme_list', [])
            if not aweme_list:
                break
            for aweme in aweme_list:
                aweme['hostname'] = hostname
                video_count += 1
                self._join_download_queue(aweme, target_folder)
                QApplication.processEvents() #刷新页面1
                #print("number: ", video_count)
                
                TextWrite(video_count)
            if contentJson.get('has_more'):
                cursor = contentJson.get('cursor')
            else:
                break
        if video_count == 0:
            TextWrite("There's no video in challenge %s." % challenge_id)
        return video_count

    def _download_music_media(self, music_id, url):
        if not music_id:
            print("Challenge #%s does not exist" % music_id)
            return
        current_folder = os.getcwd()
        target_folder = os.path.join(current_folder, 'download/@%s' % music_id)
        if not os.path.isdir(target_folder):
            os.mkdir(target_folder)

        hostname = urllib.parse.urlparse(url).hostname
        signature = self.generateSignature(str(music_id))
        music_video_url = "https://%s/aweme/v1/music/aweme/?{0}" % hostname
        music_video_params = {
            'music_id': str(music_id),
            'count': '9',
            'cursor': '0',
            'aid': '1128',
            'screen_limit': '3',
            'download_click_limit': '0',
            '_signature': signature
        }

        cursor, video_count = None, 0
        while True:
            if cursor:
                music_video_params['cursor'] = str(cursor)
                music_video_params['_signature'] = self.generateSignature(
                    str(music_id) + '9' + str(cursor))

            url = music_video_url.format(
                '&'.join([key + '=' + music_video_params[key] for key in music_video_params]))
            res = requests.get(url, headers=HEADERS)
            contentJson = json.loads(res.content.decode('utf-8'))
            aweme_list = contentJson.get('aweme_list', [])
            if not aweme_list:
                break
            for aweme in aweme_list:
                aweme['hostname'] = hostname
                video_count += 1
                self._join_download_queue(aweme, target_folder)
            if contentJson.get('has_more'):
                cursor = contentJson.get('cursor')
            else:
                break
        if video_count == 0:
            print("There's no video in music %s." % music_id)
        return video_count

def parse_sites(fileName):
    with open(fileName, "rb") as f:
        txt = f.read().rstrip().lstrip()
        txt = codecs.decode(txt, 'utf-8')
        txt = txt.replace("\t", ",").replace(
            "\r", ",").replace("\n", ",").replace(" ", ",")
        txt = txt.split(",")
    numbers = list()
    for raw_site in txt:
        site = raw_site.lstrip().rstrip()
        if site:
            numbers.append(site)
    return numbers

download_favorite = False






#界面
class MainForm(QWidget):
    def __init__(self, name = 'MainForm'):
        super(MainForm,self).__init__()
        self.setWindowTitle(name)
        self.cwd = os.getcwd() # 获取当前程序文件位置
        self.resize(479, 547)   # 设置窗体大小

        self.selectfile = QtWidgets.QPushButton(self)
        self.selectfile.setGeometry(QtCore.QRect(370, 60, 81, 23))
        self.selectfile.setObjectName("selectfile")
        self.selectfile.setText("选取文件")

        #文件路径label
        self.label_3 = QtWidgets.QLabel(self)
        self.label_3.setGeometry(QtCore.QRect(20, 60, 71, 21))
        self.label_3.setObjectName("label_3")
        self.label_3.setText("文件路径：")

        # 显示文件路径
        self.lujinuot = QtWidgets.QLabel(self)
        self.lujinuot.setGeometry(QtCore.QRect(90, 60, 271, 20))
        self.lujinuot.setObjectName("lujinuot")
        self.lujinuot.setText(" ")

        # 直接输入地址
        self.input_url = QtWidgets.QLineEdit(self)
        self.input_url.setGeometry(QtCore.QRect(90, 118, 351, 21))
        self.input_url.setObjectName("input_url")
        self.input_url.setText(" ")

        # URL标题
        self.label_url = QtWidgets.QLabel(self)
        self.label_url.setGeometry(QtCore.QRect(20, 120, 54, 12))
        self.label_url.setObjectName("label_url")
        self.label_url.setText("地址URL：")

        # 开始按钮
        self.startpachong = QtWidgets.QPushButton(self)
        self.startpachong.setGeometry(QtCore.QRect(20, 190, 431, 41))
        self.startpachong.setObjectName("startpachong")
        self.startpachong.setText("开始爬取")

        # 结果输出
        self.textout = QtWidgets.QTextBrowser(self)
        self.textout.setGeometry(QtCore.QRect(20, 251, 431, 271))
        self.textout.setObjectName("textout")

        # 设置信号槽
        self.selectfile.clicked.connect(self.slot_selectfile)
        self.textout.append("输入地址后，点开始爬取，支持以话题分享链接或用户分享链接爬取。\n")
        # 开始爬取信号
        self.startpachong.clicked.connect(StartPa)


    def slot_selectfile(self):
    #def slot_btn_chooseFile(self):
        fileName_choose, filetype = QFileDialog.getOpenFileName(self,
                                    "选取文件",
                                    self.cwd, # 起始路径
                                    "All Files (*);;Text Files (*.txt)")   # 设置文件扩展名过滤,用双分号间隔

        if fileName_choose == "":
            print("\n取消选择")
            return

        TextWrite("\n你选择的文件为:")
        TextWrite(fileName_choose)
        print("文件筛选器类型: ",filetype)
        self.lujinuot.setText(fileName_choose)

def TextWrite(msg):
    mainForm.textout.append(str(msg))
    print(msg)
    #QtWidgets.QApplication.processEvents()

def StartPa():
    if mainForm.lujinuot.text() == " " and mainForm.input_url.text() == " ":
       TextWrite("<html><head/><body><span style=\" color:red;\">参数文件或URL不能为空!!!</span></body></html>")
       return
    else:
        pa()

def pa():
    #print("pa!")
    content, opts, args = None, None, []
    try:
        #if len(sys.argv) >= 2:
        if mainForm.input_url.text() != " ":
            #opts, args = getopt.getopt(sys.argv[1:], "hi:o:", ["favorite"])
            args = [mainForm.input_url.text()]
    except getopt.GetoptError as err:
        
        sys.exit(2)

    if not args:
        # check the sites file"
        #filename = "share-url.txt"
        filename = mainForm.lujinuot.text()
        if os.path.exists(filename):
            content = parse_sites(filename)
        else:
            
            sys.exit(1)
    else:
        content = (args[0] if args else '').split(",")

    if len(content) == 0 or content[0] == "":
        
        sys.exit(1)

    if opts:
        for o, val in opts:
            if o in ("--favorite"):
                download_favorite = True
                break

    CrawlerScheduler(content)

if __name__=="__main__":
    #界面部分：
    app = QApplication(sys.argv)
    mainForm = MainForm('抖动的西蒙 v0.1')
    mainForm.show()
    sys.exit(app.exec_())
